import sys
import os
from pathlib import Path

import numpy as np
from keras.optimizers import SGD
from ACRmodel.my_model import MyModel
from ACRmodel.preprocess import preprocess_bed_to_npy
from tensorflow.keras import backend as K
import tensorflow as tf
import gc

os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'  # 0=全部信息, 1=忽略通知, 2=忽略警告, 3=忽略错误


def predict_plant(predict_id, plant_name, model_id, files):
    """
    预测前的预处理
    Args:
        predict_id: 预测记录ID
        plant_name: 植物名称
        model_id: 模型ID
        files: 上传的bed文件
    Returns:
        predict_result = {
            "status": "预测失败",  # 预测记录的状态: 未预测,预测中,预测失败,预测成功
            "chr_all": "",  # 染色体名称
            "chr_count": "",  # 染色体种类数量
            "message": "预测信息"
        }
    """
    predict_result = {
        "status": "预测失败",  # 预测记录的状态: 未预测,预测中,预测失败,预测成功
        "chr_all": "",  # 染色体名称
        "chr_count": "",  # 染色体种类数量
        "message": "服务器错误，预测失败"
    }
    try:
        # 获取项目根目录
        current_file_path = os.path.abspath(__file__)  # 预期：F:\Project\PlantACR\ACRmodel\predict.py
        project_root = os.path.dirname(os.path.dirname(current_file_path))  # 预期：F:\Project\PlantACR

        data_dir = Path(project_root) / "ACRmodel" / "static" / "predict_data" / plant_name / str(predict_id)
        data_file = Path(data_dir) / "data_predict.npy"
        save_dir = Path(project_root) / "ACRmodel" / "static" / "predict_result" / plant_name / str(predict_id)
        os.makedirs(save_dir, exist_ok=True)

        # 检查文件是否存在
        if not data_file.exists():
            print("\033[1;30;34m%s\033[0m" % "预处理数据中...")
            # 立即输出缓冲区
            sys.stdout.flush()
            result, mes, chrom_count, chrom_str, file_size_mb, sample_count = preprocess_bed_to_npy(
                plant_name, predict_id, files, 1)
            print("\033[1;30;34m%s\033[0m" % "预处理结束")
            if not result:
                predict_result["message"] = mes
                return predict_result
        else:
            print("\033[1;30;34m%s\033[0m" % "已完成预处理，跳过预处理步骤")
            sys.stdout.flush()

        # 预处理结束后，检查文件是否不存在
        if not data_file.exists():
            predict_result["message"] = "文件不存在"
            return predict_result

        # 构建与训练模型相同结构的模型
        model = MyModel()
        sgd = SGD(learning_rate=float(0.001), decay=1e-5, momentum=0.9, nesterov=True)
        model.compile(loss='binary_crossentropy', optimizer=sgd, metrics=['accuracy'])

        # 加载模型权重（也就是我们训练好的模型）
        model_weights_path = Path(project_root) / "ACRmodel" / "static" / "weights" / plant_name / str(
            model_id) / "model_weights_best"
        status = model.load_weights(model_weights_path)

        print("\033[1;30;34m%s\033[0m" % "开始预测...")
        sys.stdout.flush()
        # 使用模型预测（读取的是数据预处理中”预测处理“的处理结果）
        data_npy = np.load(Path(data_dir) / "data_predict.npy")  # 预测处理结果文件
        pred = model.predict(data_npy)
        pred_flat = np.round(pred.flatten(), 6)  # 将pred转换为一维数组
        # 阈值0.5 大于阈值为1，即有可及性; 小于阈值为0，即无可及性
        threshold = 0.5
        pred_binary = np.where(pred > threshold, 1, 0)
        pred_binary_flat = pred_binary.flatten()  # 将pred_binary转换为一维数组
        # 忽略警告
        status.expect_partial()
        print("\033[1;30;34m%s\033[0m" % "预测完成")

        # 读取bed文件内容
        bed_file = Path(data_dir) / f"{predict_id}.bed"
        with open(bed_file, 'r') as bed:
            bed_lines = bed.readlines()

        # 统计染色体信息
        chromosomes = set()
        for line in bed_lines:
            if line.strip():  # 跳过空行
                chrom = line.split('\t')[0]  # 提取染色体名称
                chromosomes.add(chrom)
        predict_result["chr_all"] = ",".join(sorted(chromosomes))  # 染色体名称列表
        predict_result["chr_count"] = str(len(chromosomes))  # 染色体种类数量

        # 融合bed文件和预测结果，创建一个bed文件
        predict_file = Path(save_dir) / "predict.bed"
        with open(predict_file, 'w') as fused_file:
            # 三部分: bed原文件内容，预测值（置信度），预测结果（0/1）
            for bed_line, pred_value, pred_binary_value in zip(bed_lines, pred_flat, pred_binary_flat):
                # 移除bed_line末尾的换行符，并用制表符连接它和预测值
                # 注意：我们不需要将pred_value转换为整数，而是直接将其格式化为字符串并保留六位小数
                fused_line = bed_line.strip() + '\t' + format(pred_value, '.6f') + '\t' + str(pred_binary_value)
                fused_file.write(fused_line + '\n')
        print("\033[1;30;34m%s\033[0m" % "生成预测文件完成")

        # # 删除temp文件夹中的文件，若要使得一条记录可以多次预测，则temp不可删除
        # if files:
        #     temp_dir = Path(project_root) / "ACRmodel" / "static" / "temp"
        #     for filename in files.split(','):
        #         file_path = temp_dir / filename.strip()
        #         if file_path.exists():
        #             os.remove(file_path)
        #     print("temp中文件删除成功")

        predict_result["status"] = "预测成功"
        predict_result["message"] = "预测成功，已生成预测文件"

        # 释放GPU内存
        K.clear_session()
        tf.compat.v1.reset_default_graph()
        gc.collect()
        return predict_result
    except Exception as e:
        print("预测失败：", e)
        predict_result["status"] = "预测失败"
        predict_result["message"] = "预测失败：" + str(e)

        # 释放GPU内存
        K.clear_session()
        tf.compat.v1.reset_default_graph()
        gc.collect()
        return predict_result


if __name__ == "__main__":
    # 测试代码
    predict_id1 = 1
    plant_name1 = "水稻"
    model_id1 = 2
    files1 = "111.bed,222.bed"
    res = predict_plant(predict_id1, plant_name1, model_id1, files1)
    print(res)
